#include <stdio.h>
#include <string.h>
#include <arpa/inet.h>

typedef struct sm3_ctx
{
    unsigned int totallen;
    unsigned int packetlen;
    unsigned int state[8];
    unsigned char packet[64];
} SM3_CTX;

// FF
#define FF1(x, y, z) ((x) ^ (y) ^ (z))
#define FF2(x, y, z) (((x) & (y)) | ((x) & (z)) | ((y) & (z)))
#define FF(j, x, y, z) ((j <= 15) ? (FF1((x), (y), (z))) : (FF2((x), (y), (z))))

// 在开发过程中 将 GG 中 > 16 是运算过程中的  ¬：32比特非运算，理解为 逻辑非!，导致错误，应当为 取反操作 ~
// GG
#define GG1(x, y, z) ((x) ^ (y) ^ (z))
#define GG2(x, y, z) (((x) & (y)) | ((~(x)) & (z)))
#define GG(j, x, y, z) ((j <= 15) ? (GG1((x), (y), (z))) : (GG2((x), (y), (z))))

// uint32_t 循环左移
#define R_LEFT(x, n) (((x) << (n)) | ((x) >> (32 - n)))

// P0 P1
#define P0(x) ((x) ^ (R_LEFT((x), 9)) ^ (R_LEFT((x), 17)))
#define P1(x) ((x) ^ (R_LEFT((x), 15)) ^ (R_LEFT((x), 23)))

#ifdef DEBUG
#define sm3prt(arg, ...)                                                                                       \
    do                                                                                                         \
    {                                                                                                          \
        fprintf(stdout, "[%s:%d]<\033[1;31m%s\033[0m> " arg, __FILE__, __LINE__, __FUNCTION__, ##__VA_ARGS__); \
    } while (0)
#else
#define sm3prt(arg, ...) ((void)0)
#endif

static const unsigned int Tj[] = {
    0x79cc4519, // 0 ≤ j ≤ 15
    0x7a879d8a, // 16 ≤ j ≤ 63
};

uint32_t big_endian_to_host(const uint8_t *data)
{
    uint32_t result;
    result = (uint32_t)data[0] << 24;
    result |= (uint32_t)data[1] << 16;
    result |= (uint32_t)data[2] << 8;
    result |= (uint32_t)data[3];
    return result;
}

int calculate_k(unsigned int l)
{
    int mask = (l + 1) & 0x1FF;
    return ((448 - mask) >= 0 ? (448 - mask) : (960 - mask));
}

uint64_t htonll(uint64_t val)
{
    return (((uint64_t)htonl(val)) << 32) + htonl(val >> 32);
}

uint64_t ntohll(uint64_t val)
{
    return (((uint64_t)ntohl(val)) << 32) + ntohl(val >> 32);
}

void sm3_init(SM3_CTX *ctx)
{
    ctx->state[0] = 0x7380166f;
    ctx->state[1] = 0x4914b2b9;
    ctx->state[2] = 0x172442d7;
    ctx->state[3] = 0xda8a0600;
    ctx->state[4] = 0xa96f30bc;
    ctx->state[5] = 0x163138aa;
    ctx->state[6] = 0xe38dee4d;
    ctx->state[7] = 0xb0fb0e4e;
    ctx->packetlen = 0;
    ctx->totallen = 0;
}

void prt_num(const char *title, const unsigned int *array, int size)
{
    if (title)
        printf("%s:\n", title);
    for (int i = 0; i < size; i++)
    {
        printf("%08x", array[i]);
        if ((i + 1) % 1 == 0)
        {
            printf(" ");
        }
        if ((i + 1) % 8 == 0)
        {
            printf("\n");
        }
    }
}

void sm3_cf(const unsigned int W[68], const unsigned int W_[64], unsigned int *state)
{
    unsigned int SS1 = 0;
    unsigned int SS2 = 0;
    unsigned int TT1 = 0;
    unsigned int TT2 = 0;

    unsigned int A, B, C, D, E, F, G, H;
    unsigned int *tmp[] = {&A, &B, &C, &D, &E, &F, &G, &H};

    for (int i = 0; i < 8; i++)
    {
        *tmp[i] = state[i];
    }

    for (int j = 0; j < 64; j++)
    {
        SS1 = R_LEFT(((R_LEFT(A, 12)) + E + (R_LEFT(Tj[(!!(j >> 4))], j))), 7);
        SS2 = SS1 ^ (R_LEFT(A, 12));
        TT1 = FF(j, A, B, C) + D + SS2 + W_[j];
        TT2 = GG(j, E, F, G) + H + SS1 + W[j];

        D = C;
        C = R_LEFT(B, 9);
        B = A;
        A = TT1;
        H = G;
        G = R_LEFT(F, 19);
        F = E;
        E = P0(TT2);
    }

    for (int i = 0; i < 8; i++)
    {
        state[i] ^= *tmp[i];
    }
    return;
}

void sm3_encode(const unsigned char *data, unsigned int W[68], unsigned int W_[64])
{
    memcpy(W, data, 64);
    for (int i = 0; i < 16; i++)
    {
        W[i] = big_endian_to_host((const unsigned char *)(data + i * 4));
    }

    for (int j = 16; j <= 67; j++)
    {
        W[j] = (P1(W[j - 16] ^ W[j - 9] ^ (R_LEFT(W[j - 3], 15)))) ^ (R_LEFT(W[j - 13], 7)) ^ W[j - 6];
    }

    for (int j = 0; j <= 63; j++)
    {
        W_[j] = W[j] ^ W[j + 4];
    }
    return;
}

void sm3_transform(const unsigned char data[64], unsigned int *state)
{
    unsigned int W[68] = {0};
    unsigned int W_[64] = {0};
    sm3_encode(data, W, W_);

    prt_num(NULL, W, 16);

    // prt_num("W", W, 68);
    // prt_num("W'", W_, 64);

    sm3_cf(W, W_, state);
    return;
}

void sm3_update(SM3_CTX *ctx, const unsigned char *data, unsigned int datalen)
{
    unsigned int partlen = 0;
    unsigned int index = 0;
    unsigned int i = 0;

    index = ctx->packetlen;
    partlen = 64 - index;

    if (datalen >= partlen)
    {
        memcpy(&ctx->packet[index], data, partlen);
        sm3_transform(ctx->packet, ctx->state);
        ctx->packetlen = 0;

        for (i = partlen; i + 64 <= datalen; i += 64)
        {
            sm3_transform(&data[i], ctx->state);
            ctx->packetlen = 0;
        }
        index = 0;
    }
    else
    {
        i = 0;
    }
    memcpy(&ctx->packet[index], &data[i], datalen - i);
    ctx->packetlen = datalen - i;
    ctx->totallen += datalen;
    sm3prt("ctx->packetlen[%u], datalen-i[%u], i[%u] datalen[%u]\n", ctx->packetlen, (datalen - i), i, datalen);
}

void sm3_padding(unsigned char *data, int len, unsigned long long ctx_len, unsigned char *pda, int *flag)
{
    char m_[128] = {0};
    unsigned int l = len * 8;
    int k = calculate_k(l);
    unsigned int j = len;
    unsigned long long lh = htonll(ctx_len * 8);

    sm3prt("l[%u], k[%u], j[%u], lh[%llu] ctx_len[%llu]\n", l, k, j, lh, ctx_len);

    memcpy(m_, data, len);

    if (k > 0)
    {
        m_[j++] = 1u << 7;
        k -= 8;
    }

    while (k > 0)
    {
        m_[j++] = 0;
        k -= 8;
    }

    sm3prt("k[%u], j[%u]\n%s\n", k, j, m_);

    memcpy(m_ + j, (char *)&lh, 8);

    if (j > 56)
    {
        *flag = 1;
        memcpy(pda, (char *)m_, 64);
        memcpy(data, (char *)m_ + 64, 64);
    }
    else
    {
        memcpy(data, (char *)m_, 64);
    }

    return;
}

void sm3_final(SM3_CTX *ctx)
{
    unsigned char pda[64] = {0};
    int flag = 0;

    sm3_padding(ctx->packet, ctx->packetlen, ctx->totallen, pda, &flag);

    if (flag)
    {
        sm3_transform(pda, ctx->state);
    }

    sm3_transform(ctx->packet, ctx->state);
    return;
}

void sm3(const char *data)
{
    char hash[68] = {0};
    SM3_CTX ctx;
    sm3_init(&ctx);
    sm3_update(&ctx, (const unsigned char *)data, strlen(data));
    sm3_final(&ctx);
    prt_num("Hash", ctx.state, 8);
    printf("=======================================================================\n");
    for (int i = 0; i < 8; i++)
    {
        snprintf(hash + i * 8, 9, "%08X", ctx.state[i]);
    }
    printf("hash: %s\n", hash);
}

int test_main()
{
    char data[] = "abc";
    const char m2[] = "abcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcd";
    const char m[] = "abcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcd";
    const char mm[] = "abcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcd";

    sm3(data);
    sm3(m2);
    sm3(m);
    sm3(mm);

    return 0;
}